import time
import logging
import os
from functools import wraps
from typing import Callable, Any, Dict, Optional
import json

try:
    import psutil
    has_psutil = True
except ImportError:
    has_psutil = False
    print("警告：psutil未安装，系统资源监控将不可用")

# 设置日志
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)

logger = logging.getLogger('kg_retrieval_monitor')

# 创建性能日志文件处理器
try:
    if not os.path.exists('logs'):
        os.makedirs('logs')
    perf_handler = logging.FileHandler('logs/performance.log', encoding='utf-8')
    perf_handler.setLevel(logging.INFO)
    perf_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    perf_handler.setFormatter(perf_formatter)
    logger.addHandler(perf_handler)
except Exception as e:
    print(f"无法创建日志文件：{e}")


class PerformanceMonitor:
    """性能监控工具"""
    
    def __init__(self):
        self.enabled = True
        # 记录所有函数执行统计数据
        self.stats = {}
    
    def measure(self, func=None, name: str = None):
        """测量函数执行时间的装饰器"""
        def decorator(fn):
            func_name = name or fn.__name__
            
            @wraps(fn)
            async def async_wrapper(*args, **kwargs):
                if not self.enabled:
                    return await fn(*args, **kwargs)
                
                start_time = time.time()
                mem_before = self._get_memory_usage()
                
                try:
                    result = await fn(*args, **kwargs)
                    status = "success"
                except Exception as e:
                    status = f"error: {str(e)}"
                    raise
                finally:
                    duration = time.time() - start_time
                    mem_after = self._get_memory_usage()
                    mem_diff = mem_after - mem_before if mem_before is not None and mem_after is not None else None
                    
                    # 更新函数统计
                    if func_name not in self.stats:
                        self.stats[func_name] = {
                            "calls": 0,
                            "total_time": 0,
                            "min_time": float("inf"),
                            "max_time": 0,
                            "errors": 0
                        }
                    
                    self.stats[func_name]["calls"] += 1
                    self.stats[func_name]["total_time"] += duration
                    self.stats[func_name]["min_time"] = min(self.stats[func_name]["min_time"], duration)
                    self.stats[func_name]["max_time"] = max(self.stats[func_name]["max_time"], duration)
                    
                    if "error" in status:
                        self.stats[func_name]["errors"] += 1
                    
                    # 记录日志
                    log_data = {
                        "function": func_name,
                        "duration_ms": round(duration * 1000, 2),
                        "status": status,
                        "memory_diff_mb": round(mem_diff / (1024 * 1024), 2) if mem_diff is not None else None,
                        "timestamp": time.time()
                    }
                    logger.info(json.dumps(log_data))
                
                return result
            
            @wraps(fn)
            def sync_wrapper(*args, **kwargs):
                if not self.enabled:
                    return fn(*args, **kwargs)
                
                start_time = time.time()
                mem_before = self._get_memory_usage()
                
                try:
                    result = fn(*args, **kwargs)
                    status = "success"
                except Exception as e:
                    status = f"error: {str(e)}"
                    raise
                finally:
                    duration = time.time() - start_time
                    mem_after = self._get_memory_usage()
                    mem_diff = mem_after - mem_before if mem_before is not None and mem_after is not None else None
                    
                    # 更新函数统计
                    if func_name not in self.stats:
                        self.stats[func_name] = {
                            "calls": 0,
                            "total_time": 0,
                            "min_time": float("inf"),
                            "max_time": 0,
                            "errors": 0
                        }
                    
                    self.stats[func_name]["calls"] += 1
                    self.stats[func_name]["total_time"] += duration
                    self.stats[func_name]["min_time"] = min(self.stats[func_name]["min_time"], duration)
                    self.stats[func_name]["max_time"] = max(self.stats[func_name]["max_time"], duration)
                    
                    if "error" in status:
                        self.stats[func_name]["errors"] += 1
                    
                    # 记录日志
                    log_data = {
                        "function": func_name,
                        "duration_ms": round(duration * 1000, 2),
                        "status": status,
                        "memory_diff_mb": round(mem_diff / (1024 * 1024), 2) if mem_diff is not None else None,
                        "timestamp": time.time()
                    }
                    logger.info(json.dumps(log_data))
                
                return result
            
            # 根据函数是否为异步函数选择对应的包装器
            import inspect
            if inspect.iscoroutinefunction(fn):
                return async_wrapper
            return sync_wrapper
        
        if func is None:
            return decorator
        return decorator(func)
    
    def _get_memory_usage(self) -> Optional[float]:
        """获取当前进程的内存使用情况"""
        if not has_psutil:
            return None
        
        try:
            process = psutil.Process(os.getpid())
            return process.memory_info().rss  # 物理内存使用量，单位字节
        except:
            return None
    
    def get_stats(self) -> Dict[str, Any]:
        """获取性能统计数据"""
        # 计算每个函数的平均执行时间
        result = {}
        for func_name, stats in self.stats.items():
            avg_time = stats["total_time"] / stats["calls"] if stats["calls"] > 0 else 0
            
            result[func_name] = {
                "calls": stats["calls"],
                "total_time_s": round(stats["total_time"], 3),
                "avg_time_ms": round(avg_time * 1000, 2),
                "min_time_ms": round(stats["min_time"] * 1000, 2) if stats["min_time"] != float("inf") else None,
                "max_time_ms": round(stats["max_time"] * 1000, 2),
                "errors": stats["errors"],
                "error_rate": round(stats["errors"] / stats["calls"] * 100, 2) if stats["calls"] > 0 else 0
            }
        
        # 添加系统资源信息
        if has_psutil:
            try:
                result["system"] = {
                    "cpu_percent": psutil.cpu_percent(interval=0.1),
                    "memory_percent": psutil.virtual_memory().percent,
                    "memory_available_mb": round(psutil.virtual_memory().available / (1024 * 1024), 2)
                }
            except:
                pass
        
        return result
    
    def reset_stats(self):
        """重置性能统计数据"""
        self.stats = {}

# 全局性能监控实例
performance_monitor = PerformanceMonitor() 